# -*- coding:utf-8 -*-
"""
双色球号码01序列格式化工具
功能: 加载原始爬虫数据 → 转换为r1-r33 01序列 → 保存到CSV(兼容特征工程输入)
支持train/predict模式分离, 分别处理全量/最近89期数据
"""
import os
import traceback
import pandas as pd
from typing import Tuple, List, Dict
import csv
# 从config导入全局配置和日志(避免重复定义)
from config.ssq_config import init_global_logger, logger, SSQ_CONFIG


class Lottery01Format:
    """彩票01序列格式化类: 处理原始数据→生成01编码→保存结果(支持训练/预测分离)"""
    def __init__(self, run_type: str):
        """初始化01序列化: 加载配置、创建目录、根据数据类型确定路径"""
        # 校验数据类型合法性
        if run_type not in ["train", "predict"]:
            raise ValueError(f"数据类型错误！仅支持 'train' 或 'predict', 当前输入: {run_type}")
        
        self.run_type = run_type  # 核心: 记录数据用途(train/predict)
        # 数据存储容器
        self.historical_data: Dict = {}  # {期号: {'date': 日期, 'red_balls': 红球列表, 'blue_ball': 蓝球}}
        self.formatted_data: List[List] = []  # 01格式化后的数据
        # 配置参数(从config统一获取)
        self.ssq_config = SSQ_CONFIG  # 双色球专属配置
        self.red_ball_count = SSQ_CONFIG['max_red_ball']

        # 新增: 根据数据类型确定保存文件名(从配置文件读取, 避免硬编码)
        self.data_folder = self.ssq_config["train_file"]["DATA_FOLDER"]
        self.raw_data_path = os.path.join(self.data_folder,self.ssq_config["train_file"]["csv_file"])
        self.output_file = os.path.join(self.data_folder,self.ssq_config["train_file"]["csv01_file"])
        if self.run_type == "predict":
            self.data_folder = self.ssq_config["predict_file"]["DATA_FOLDER"]
            self.raw_data_path = os.path.join(self.data_folder,self.ssq_config["predict_file"]["csv_file"])
            self.output_file = os.path.join(self.data_folder,self.ssq_config["predict_file"]["csv01_file"])

        # 创建数据存储目录(不存在则自动创建)
        if not os.path.exists(self.data_folder):
            os.makedirs(self.data_folder)
            logger.info(f"创建数据存储目录: {os.path.abspath(self.data_folder)}")
        
        # 按数据类型区分日志文件
        init_global_logger(log_file=f"format01_{run_type}_data.log", pid=os.getpid())
        logger.info(f"01序列格式化工具初始化完成(数据用途: {run_type})")
        logger.debug(f"原始数据路径: {self.raw_data_path}")
        logger.debug(f"01格式化输出路径: {self.output_file}")


    def load_data(self) -> Tuple[Dict, List[int]]:
        """
        加载对应类型的原始双色球数据(无表头CSV)
        Returns:
            Tuple[历史数据字典, 排序后期号列表]
        """
        try:
            logger.info(f"开始加载{self.run_type}类型原始数据: {os.path.abspath(self.raw_data_path)}")
            # 读取无表头CSV(格式: 期号,日期,红球1-6,蓝球)
            df = pd.read_csv(
                self.raw_data_path,
                header=None,  # 明确无表头
                dtype=str,    # 先按字符串读取, 避免数字格式丢失
                encoding="utf-8"
            )

            # 校验列数(需至少9列: 期号+日期+6红球+1蓝球)
            if df.shape[1] < 9:
                raise ValueError(
                    f"原始CSV格式错误: 需9列(期号+日期+6红球+1蓝球), 实际{df.shape[1]}列"
                )

            # 逐行解析数据
            for row_idx, row in df.iterrows():
                try:
                    if row_idx == 0: continue  # 跳过可能的表头行
                    # 1. 提取核心字段(按爬虫输出的列顺序)
                    period = int(row[0])  # 第0列: 期号
                    date_str = str(row[1]).strip()  # 第1列: 日期(去空格)
                    red_balls = [int(row[col]) for col in range(2, 8) if pd.notna(row[col])]  # 第2-7列: 红球
                    blue_ball = int(row[8])  # 第8列: 蓝球

                    # 2. 数据合法性校验
                    if len(red_balls) != 6:
                        logger.warning(f"第{row_idx+1}行数据无效: 红球数量{len(red_balls)}个(需6个), 已跳过")
                        continue
                    for red in red_balls:
                        if not (1 <= red <= self.red_ball_count):
                            logger.warning(f"第{row_idx+1}行数据无效: 红球{red}超出范围(1-33), 已跳过")
                            raise ValueError(f"红球{red}非法")
                    if not (1 <= blue_ball <= 16):
                        logger.warning(f"第{row_idx+1}行数据无效: 蓝球{blue_ball}超出范围(1-16), 已跳过")
                        continue

                    # 3. 存储到历史数据字典(红球排序)
                    self.historical_data[period] = {
                        "date": date_str,
                        "red_balls": sorted(red_balls),
                        "blue_ball": blue_ball
                    }

                except (ValueError, TypeError) as e:
                    logger.warning(f"第{row_idx+1}行数据解析失败: {str(e)}, 已跳过")
                    continue

            # 4. 生成排序后的期号列表(确保时序正确)
            sorted_periods = sorted(self.historical_data.keys())
            if not sorted_periods:
                logger.warning(f"未加载到任何{self.run_type}类型有效数据")
                return {}, []

            # 5. 输出加载结果日志
            logger.info(f"{self.run_type}类型原始数据加载完成: 共{len(sorted_periods)}期有效数据")
            logger.info(f"期号范围: {sorted_periods[0]} ~ {sorted_periods[-1]}")
            # 打印最早/最新一期数据(验证加载正确性)
            earliest_data = self.historical_data[sorted_periods[0]]
            latest_data = self.historical_data[sorted_periods[-1]]
            logger.info(
                f"最早一期: 期号{sorted_periods[0]} | 日期{earliest_data['date']} | "
                f"红球{earliest_data['red_balls']} | 蓝球{earliest_data['blue_ball']}"
            )
            logger.info(
                f"最新一期: 期号{sorted_periods[-1]} | 日期{latest_data['date']} | "
                f"红球{latest_data['red_balls']} | 蓝球{latest_data['blue_ball']}"
            )

            return self.historical_data, sorted_periods

        except FileNotFoundError:
            logger.error(f"{self.run_type}类型原始数据文件未找到: {os.path.abspath(self.raw_data_path)}")
            raise
        except Exception as e:
            logger.error(f"{self.run_type}类型原始数据加载失败: {str(e)},{traceback.format_exc()}")
            raise

    def format01_red_balls(self) -> List[List]:
        """
        将原始红球转换为01序列(核心方法)
        输出格式: [idx(期号), date(日期), r1, r2, ..., r33, b1(蓝球)]
        其中r1-r33: 1=号码出现, 0=号码未出现
        """
        self.formatted_data = []
        if not self.historical_data:
            logger.warning("无历史数据可格式化！请先调用load_data加载原始数据")
            return []

        # 逐期生成01序列
        for period, data in self.historical_data.items():
            try:
                # 1. 初始化行数据(期号+日期)
                row = [period, data["date"]]

                # 2. 生成r1-r33的01编码(按号码1-33顺序)
                red_set = set(data["red_balls"])
                for ball_num in range(1, self.red_ball_count + 1):
                    row.append(1 if ball_num in red_set else 0)

                # 3. 添加蓝球
                row.append(data["blue_ball"])

                # 4. 添加到格式化列表
                self.formatted_data.append(row)

            except Exception as e:
                logger.warning(f"期号{period}数据格式化失败: {str(e)}, 已跳过")
                continue

        # 5. 按期号排序(确保时序正确)
        self.formatted_data.sort(key=lambda x: x[0])

        # 6. 输出格式化结果日志
        logger.info(f"{self.run_type}类型01序列格式化完成: 共{len(self.formatted_data)}期有效数据")
        logger.info(f"01序列结构: {len(self.formatted_data[0])}列(期号+日期+33个01编码+蓝球)")
        return self.formatted_data

    def save_to_csv(self) -> bool:
        """保存01格式化数据到对应类型的CSV文件"""
        if not self.formatted_data:
            logger.warning("无格式化数据可保存！请先调用format01_red_balls生成01序列")
            return False

        try:
            # 先删除旧文件(避免数据覆盖冲突)
            if os.path.exists(self.output_file):
                os.remove(self.output_file)
                logger.debug(f"已删除旧{self.run_type}类型01数据文件")

            # 校验表头与数据列数一致性
            csv01_header = self.ssq_config["csv01_header"]
            if len(self.formatted_data[0]) != len(csv01_header):
                raise ValueError(
                    f"数据列数与表头不匹配！数据{len(self.formatted_data[0])}列, 表头{len(csv01_header)}列"
                )

            # 生成DataFrame并保存
            df = pd.DataFrame(
                self.formatted_data,
                columns=csv01_header  # 使用config统一表头, 兼容特征工程
            )
            df.to_csv(
                self.output_file,
                index=False,  # 不保留行索引
                encoding="utf-8",
                quoting=csv.QUOTE_MINIMAL
            )

            # 输出保存日志
            logger.info(f"{self.run_type}类型01格式化数据已保存到: {os.path.abspath(self.output_file)}")
            logger.info(
                f"保存详情: {len(df)}行数据, {len(df.columns)}列(表头: {','.join(csv01_header[:5])}...{','.join(csv01_header[-3:])})"
            )
            return True

        except Exception as e:
            logger.error(f"{self.run_type}类型01数据保存失败: {str(e)},{traceback.format_exc()}")
            return False

    def run(self) -> None:
        """完整流程: 加载对应类型原始数据 → 生成01序列 → 保存到对应CSV"""
        logger.info("=" * 80)
        logger.info(f"          双色球{self.run_type}类型01序列格式化工具          ")
        logger.info("=" * 80)

        try:
            # 1. 加载原始数据
            logger.info(f"步骤[1/3]: 加载{self.run_type}类型原始数据")
            _, _ = self.load_data()

            # 2. 生成01序列
            logger.info(f"步骤[2/3]: 生成r1-r33 01序列...")
            self.format01_red_balls()

            # 3. 保存到CSV
            logger.info(f"步骤[3/3]: 保存{self.run_type}类型01格式化数据...")
            save_ok = self.save_to_csv()
            if not save_ok:
                raise RuntimeError(f"{self.run_type}类型01数据保存失败, 任务终止")

            # 4. 打印样例数据(验证正确性)
            logger.info(f"[{self.run_type}类型格式化结果样例](最新2期数据)")
            header = self.ssq_config["csv01_header"]
            logger.info(f"表头(关键列): {','.join(header[:3])}...{','.join(header[-3:])}")
            for row in self.formatted_data[-2:]:  # 显示最新2期
                display_row = [str(row[0]), row[1]] + [str(x) for x in row[2:7]] + [str(row[-1])]
                logger.info(f"数据: {','.join(display_row)}")

            logger.info("=" * 80)
            logger.info(f"          {self.run_type}类型01序列格式化任务全部完成！          ")
            logger.info("=" * 80)
            return True
        except FileNotFoundError:
            logger.critical(
                f"致命错误: {self.run_type}类型原始数据文件不存在 → {self.raw_data_path}\n"
                f"请先运行spider脚本获取{self.run_type}类型原始数据"
            )
            return True
        except Exception as e:
            logger.critical(f"致命错误: {self.run_type}类型格式化任务异常终止 → {str(e)},{traceback.format_exc()}")
            return False


# ---------------------- 脚本入口(支持命令行参数) ----------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="双色球01序列格式化工具(支持训练/预测分离)")
    parser.add_argument(
        "--run-type",
        required=True,
        choices=["train", "predict"],
        help="数据用途: train(训练数据, 处理全量)、predict(预测数据, 处理最近89期)"
    )
    args = parser.parse_args()

    try:
        # 初始化并运行对应类型的格式化工具
        format_tool = Lottery01Format(run_type=args.run_type)
        format_tool.run()
    except Exception as e:
        logger.critical(f"脚本运行失败: {str(e)},{traceback.format_exc()}")
        exit(1)